import os

import numpy as np

'''

下载的测试集是没有切片的，
使用这段代码对数据集进行切片，
并且去除只有背景的图像
'''


def process_file(file_path):
    try:
        # 加载图像数据
        data = np.load(file_path)

        # 假设标签文件与图像文件在对应的目录，并具有相同的文件名
        label_path = file_path.replace("images", "labels")
        if not os.path.exists(label_path):
            print(f"标签文件不存在: {label_path}")
            return

        # 加载标签数据
        label_data = np.load(label_path)

        # 数据预处理（根据需要进行修改）
        # 例如，如果数据值在 -125 到 275 之间，可以进行裁剪和归一化
        # data_clipped = np.clip(data, -125, 275)
        # data_normalized = (data_clipped - (-125)) / (275 - (-125))
        # del data_clipped

        # 确保输出目录存在
        output_dir = "./processed_slices"
        os.makedirs(output_dir, exist_ok=True)

        # 对数据进行切片处理（假设沿第三个维度切片）
        for i in range(data.shape[0]):
            formatted_i = "{:04d}".format(i + 1)

            # 获取当前切片的数据和标签
            slice_data = data[i, :, :]
            label_slice = label_data[i, :, :]

            # 验证标签值（根据实际情况调整）
            unique_labels = np.unique(label_slice)
            if len(unique_labels) <= 4:
                continue
            print(unique_labels)
            if not all(label in [0, 1, 2, 3, 4, 5, 6, 7, 8] for label in unique_labels):
                print(f"警告：在切片 {i} 中发现意外的标签值: {unique_labels}")

            # 生成保存文件名
            case_name = os.path.splitext(os.path.basename(file_path))[0]
            slice_filename = f"{case_name}_slice_{formatted_i}.npy"
            label_filename = f"{case_name}_slice_{formatted_i}_label.npy"

            slice_file_path = os.path.join(output_dir + '/images', slice_filename)
            label_file_path = os.path.join(output_dir + '/labels', label_filename)
            # 保存切片数据
            np.save(os.path.join(output_dir + '/images', slice_filename), slice_data)
            np.save(os.path.join(output_dir + '/labels', label_filename), label_slice)
            with open('test_dir.txt', 'a') as f:
                f.write(slice_file_path + ' ' + label_file_path + '\n')



    except Exception as e:
        print(f"处理文件 {file_path} 时发生错误: {str(e)}")


def main():
    data_path = "/Volumes/For_Mac/dateset/Synapse_npy/test/images"

    for root, dirs, files in os.walk(data_path):
        for file in files:
            if file.endswith(".npy"):
                file_path = os.path.join(root, file)
                process_file(file_path)


if __name__ == "__main__":
    main()
